# coding: utf8

"""
基于pandas导入导出excel相关
"""

# # 全局函数 - DataFrame导出的Excel列宽自适应
# # Python==3.6.4、Pandas==1.0.1、numpy==1.19.5、openpyxl==3.0.9
from pathlib import Path

import numpy as np
import pandas as pd
from openpyxl.utils import get_column_letter


def to_excel_auto_column_weight(
        df: pd.DataFrame,
        file_path=r"./df_to_excel.xlsx",
        sheet_name="Sheet1",
        header=True) -> str:
    """
    DataFrame导出为Excel并自动设置列宽

    :param df: pd.DataFrame数据表
    :param file_path: excel文件保存路径
    :param sheet_name: excel工作簿名称，默认Sheet1
    :param header: 是否导出df列头做为表格列名
    :return: 保存后的excel表格文件路径
    """
    if not isinstance(df, pd.DataFrame):
        raise ValueError(f"df参数需要pd.DataFrame数据类型，而非{type(df)}。")

    # 保存路径目录不存在则创建
    file_path = str(file_path)
    if not Path(file_path).parent.exists():
        Path(file_path).parent.mkdir(parents=True, exist_ok=True)

    with pd.ExcelWriter(file_path) as writer:
        # 数据 to 写入器，并指定sheet名称
        df.to_excel(writer, sheet_name=sheet_name,
                    index=False, header=header)

        #  计算每列表头的字符宽度
        column_widths = (
            df.columns.to_series().apply(
                lambda x: len(str(x).encode('gbk'))
            ).values
        )

        #  计算每列的最大字符宽度
        max_widths = (
            df.astype(str).applymap(
                lambda x: len(str(x).encode('gbk'))
            ).agg(max).values
        )

        # 取前两者中每列的最大宽度
        widths = np.max([column_widths, max_widths], axis=0)

        # 指定sheet，设置该sheet的每列列宽
        worksheet = writer.sheets[sheet_name]
        for i, width in enumerate(widths, 1):
            # openpyxl引擎设置字符宽度时会缩水0.5左右个字符，
            # 所以干脆+2使左右都空出一个字宽。
            worksheet.column_dimensions[
                get_column_letter(i)].width = width + 2
    return Path(file_path).absolute().as_posix()


# # 导出excel时，列设置货币和其他数值类型
# # Python==3.8.10、Pandas==1.4.4、openpyxl==3.0.10
import pandas as pd


def to_excel_column_currency_format(
        df: pd.DataFrame,
        columns: [str, list] = None,
        file_path=r"./df_to_excel_currency_format.xlsx",
        sheet_name="Sheet1",
        header=True) -> str:
    """
    导出excel时设置一列或多列为货币数值格式

    :param df: pd.DataFrame数据表
    :param columns: 某列：'列名' 字符串，多列：["列1", "列2"] 列表
    :param file_path: excel文件保存路径
    :param sheet_name: excel工作簿名称，默认Sheet1
    :param header: 是否导出df列头做为表格列名
    :return: 保存后的excel表格文件路径
    """
    if not isinstance(df, pd.DataFrame):
        raise ValueError(f"df参数需要pd.DataFrame数据类型，而非{type(df)}。")

    # 保存路径目录不存在则创建
    file_path = str(file_path)
    if not Path(file_path).parent.exists():
        Path(file_path).parent.mkdir(parents=True, exist_ok=True)

    # 空表格或者没有列设置，则直接导出excel表格
    if columns is None or df.empty:
        df.to_excel(file_path, sheet_name=sheet_name, index=False, header=header)
        return Path(file_path).absolute().as_posix()

    if not isinstance(columns, list):
        columns = [columns]

    # 数值列处理成浮点数（含有逗号的字符串货币）
    excel_columns = []
    for column in columns:
        df[column] = df[column].astype(str).map(
            lambda s: s.replace(",", "")).astype(float)
        excel_columns.append(df.columns.get_loc(column))

    # 货币格式，导出excel
    with pd.ExcelWriter(file_path) as writer:
        # 数据 to 写入器，并指定sheet名称
        df.to_excel(writer, sheet_name=sheet_name,
                    index=False, header=header)

        # 设置excel列数值格式，其他格式参考一下链接
        # https://xlsxwriter.readthedocs.io/format.html#set_num_format
        workbook = writer.book
        cell_format = workbook.add_format({'num_format': '###0.00'})
        worksheet = writer.sheets[sheet_name]

        # 逐列设置格式
        for excel_column in excel_columns:
            first_col, last_col = excel_column, excel_column + 1
            worksheet.set_column(
                first_col=first_col, last_col=last_col, cell_format=cell_format)
    return Path(file_path).absolute().as_posix()


if __name__ == "__main__":
    # 自适应宽度
    p = to_excel_auto_column_weight(pd.DataFrame({0: [1, 2, 3], "a": [4, 5, 6]}))
    print(p)

    # 货币数值格式
    p = to_excel_column_currency_format(pd.DataFrame({
        "name": ["总额", "单价", "数量", "原价", "损耗"],
        "amount": [2000000000.245, 19.2, '28,000', 27.56789, 24],
    }), "amount")
    print(p)
